from sentence_transformers import SentenceTransformer
import faiss
import numpy as np

def create_index_and_embeddings(documents):
    """
    创建 FAISS 向量数据库索引和文档嵌入向量。

    Args:
        documents: 文档列表 (list of str)

    Returns:
        tuple: (FAISS index, document embeddings)
    """
    # 加载预训练的 SentenceTransformer 嵌入模型
    embedding_model = SentenceTransformer('all-MiniLM-L12-v2')

    # 创建文档嵌入向量
    doc_embeddings = embedding_model.encode(documents)

    # 使用 FAISS 构建向量数据库索引
    embedding_dimension = doc_embeddings.shape[1] # 获取嵌入维度
    index = faiss.IndexFlatL2(embedding_dimension) # 使用 L2 距离作为相似度度量
    index.add(np.array(doc_embeddings))  # 将文档嵌入向量添加到索引中

    print("向量数据库索引构建完成！")
    return index, doc_embeddings

if __name__ == '__main__':
    # 模拟企业文档 (实际应用中应从文件或数据库加载)
    documents = [
        "我们的退货政策是客户可以在30天内退货。",
        "2025年新的假期安排，每人享有法定年假12天。",
        "上个月，我们的销售政策是全场商品8折优惠。",
        "公司鼓励员工参与志愿活动，每年提供一天带薪志愿假。"
    ]

    index, doc_embeddings = create_index_and_embeddings(documents)

    # 可以保存 index 和 embeddings 到文件，以便后续加载使用
    # faiss.write_index(index, "document_index.faiss")
    # np.save("document_embeddings.npy", doc_embeddings)

    print("\n--- 索引创建和存储示例完成 ---")